package main

import (
	"context"
	"encoding/hex"
	"flag"
	"fmt"
	"golang.org/x/sync/semaphore"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

var checksumContent string

// calculateHash 计算文件的哈希值
func calculateHash(filePath string, algoI HashAlgorithmInterface) (string, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return "", err
	}
	defer file.Close()

	// 获取文件信息用于计算进度
	fileInfo, err := file.Stat()
	if err != nil {
		return "", err
	}
	totalSize := fileInfo.Size()
	processedSize := int64(0)

	hasher := algoI.New()
	buf := make([]byte, 1024*1024)
	for {
		n, err := file.Read(buf)
		if err != nil && err != io.EOF {
			return "", err
		}
		if n == 0 {
			break
		}

		hasher.Write(buf[:n])

		// 计算进度
		processedSize += int64(n)
		progress := float64(processedSize) / float64(totalSize) * 100
		progressBarLength := 50
		filledLength := int(progress / 100 * float64(progressBarLength))
		progressBar := fmt.Sprintf("[%s%s] %.2f%%\r", strings.Repeat("#",filledLength), strings.Repeat(" ", progressBarLength-filledLength), progress)
		//fmt.Printf("Progress: %.2f%%\r", progress)
		fmt.Print(progressBar)
	}

	return hex.EncodeToString(hasher.Sum(nil)), nil
}

// processFile 并发处理单个文件
func processFile(path string, wg *sync.WaitGroup, sem *semaphore.Weighted, processedFiles *int32, algo HashAlgorithmInterface) {
	defer wg.Done()
	defer sem.Release(1)

	hash, err := calculateHash(path, algo)
	if err != nil {
		fmt.Printf("Error processing file %s: %v\n", path, err)
		return
	}

	fmt.Printf("File: %s\nHash: %s\n\n", path, hash)
	checksumContent += path + "\n " + hash + "\n\n"
	atomic.AddInt32(processedFiles, 1)
}

// processDirectory 遍历目录并计算每个文件的哈希值
func processDirectory(dirPath string, maxConcurrency int, algo HashAlgorithmInterface) (int32, error) {
	var wg sync.WaitGroup
	sem := semaphore.NewWeighted(int64(maxConcurrency))
	var processedFiles int32

	err := filepath.Walk(dirPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if !info.IsDir() {
			wg.Add(1)
			if err := sem.Acquire(context.Background(), 1); err != nil {
				return err
			}
			go processFile(path, &wg, sem, &processedFiles, algo)
		}
		return nil
	})

	wg.Wait()
	return processedFiles, err
}

func main() {
	maxConcurrency := 8

	dirPath := flag.String("d", "", "需要指定的扫描路径")
	hashAlgo := flag.String("t", "", "需要指定hash类型: sha256, crc32, md5, blake2b, crc64")
	flag.Parse()

	if *hashAlgo == "" {
		*hashAlgo = "crc64"
	}

	if *dirPath == "" {
		fmt.Println("Usage:")
		flag.PrintDefaults()
		//*dirPath = "./test/test-directory"
		os.Exit(1)
	}

	algo := hashAlgoMap[*hashAlgo]

	// 计算Execution time
	mainStart := time.Now()
	processedFiles, err := processDirectory(*dirPath, maxConcurrency, algo)
	if err != nil {
		fmt.Printf("Error processing directory: %v\n", err)
		return
	}

	fmt.Printf("Total processed files: %d\n", processedFiles)
	mainElapsed := time.Since(mainStart)
	fmt.Printf("Execution time: %s\n", mainElapsed)
	errWrite := writeFile("./checksum", checksumContent)
	if errWrite != nil {
		fmt.Println(errWrite)
		return
	}
}

func writeFile(filePath string, content string) error {
	file, err := os.OpenFile(filePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return fmt.Errorf("failed to write to file: %v", err)
	}
	defer file.Close()
	_, err = file.WriteString(content)
	if err != nil {
		return fmt.Errorf("failed to write to file: %v", err)
	}

	return nil
}
